'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: dalle3.py
* @Time: 2025/9/22
* @All Rights Reserve By Brtc
'''
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from internal.lib.helper import add_attribute


class Dalle3ArgsSchema(BaseModel):
    query:str = Field(description="输入是生成图片的提示词(prompt)")

@add_attribute("args_schema", Dalle3ArgsSchema)
def dalle3(**kwargs)->BaseTool:
    """返回dalle3绘图的langchain 工具"""
    return OpenAIDALLEImageGenerationTool(
        api_wrapper=DallEAPIWrapper(model="dall-e-3", **kwargs),
        args_schema=Dalle3ArgsSchema,
    )